Skip to content

mxfp8 training: add TP sharding strategy for dim1 kernel #2436

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 25 commits into from
Jul 9, 2025
Merged

mxfp8 training: add TP sharding strategy for dim1 kernel #2436

merged 25 commits into from
Jul 9, 2025

Conversation

vkuzo
Copy link
Contributor

@vkuzo vkuzo commented Jun 24, 2025 •
edited
Loading

Summary:

Enables mxfp8 training with the dim1 triton kernel and TP, in eager mode. In detail:

  1. write a sharding strategy for the triton kernel custom op
  2. properly wrap the outputs of the custom op in DTensor
  3. add a test which enables the triton custom op, and modify the test model sizes since this kernel needs block_size 32

Note that compile does not work yet, seeing https://www.internalfb.com/phabricator/paste/view/P1851219639

Test Plan:

// TP test passes (with eager mode)
./test/prototype/mx_formats/test_mx_dtensor.sh

// unit tests still pass
pytest test/prototype/mx_formats

// torchtitan debug model loss goes down (in eager)
with-proxy CONFIG_FILE="torchtitan/models/llama3/train_configs/debug_model.toml " ./run_train.sh --model.print_after_conversion --training.steps 50 --model.converters mx --mx.recipe_name "mxfp8" --parallelism.tensor_parallel_degree=2 --mx.use_fp8_dim1_cast_triton_kernel

Reviewers:

Subscribers:

Tasks:

Tags:

vkuzo added 22 commits June 20, 2025 07:10
@vkuzo
Update
5c23c6b 
[ghstack-poisoned]
@vkuzo
Update
ad2ce62 
[ghstack-poisoned]
@vkuzo
Update
5eb2066 
[ghstack-poisoned]
@vkuzo
Update
6e3df57 
[ghstack-poisoned]
@vkuzo
Update
75e6fe7 
[ghstack-poisoned]
@vkuzo
Update
8bf42da 
[ghstack-poisoned]
@vkuzo
Update
c0080cd 
[ghstack-poisoned]
@vkuzo
Update
c6fc48b 
[ghstack-poisoned]
@vkuzo
Update
4cc1531 
[ghstack-poisoned]
@vkuzo
Update
42083e2 
[ghstack-poisoned]
@vkuzo
Update
9d171ad 
[ghstack-poisoned]
@vkuzo
Update
09c1c58 
[ghstack-poisoned]
@vkuzo
Update
e511e7b 
[ghstack-poisoned]
@vkuzo
Update
3562a5e 
[ghstack-poisoned]
@vkuzo
Update
7a0fd96 
[ghstack-poisoned]
@vkuzo
Update
2d1545f 
[ghstack-poisoned]
@vkuzo
Update
20b7db2 
[ghstack-poisoned]
@vkuzo
Update
7788412 
[ghstack-poisoned]
@vkuzo
Update
aabeb61 
[ghstack-poisoned]
@vkuzo
Update
28f32b9 
[ghstack-poisoned]
@vkuzo
Update
1001602 
[ghstack-poisoned]
@vkuzo
Update
b251829 
[ghstack-poisoned]
@vkuzo
Copy link
Contributor Author

vkuzo commented Jun 24, 2025 •
edited
Loading

Stack from ghstack (oldest at bottom):

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jun 24, 2025 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2436

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Cancelled Job

As of commit 68b683d with merge base c57226b (image):

NEW FAILURE - The following job has failed:

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@vkuzo vkuzo mentioned this pull request Jun 24, 2025
Merged
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 24, 2025
vkuzo added a commit that referenced this pull request Jun 24, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
3737967 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 3b7f85d
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo
Update
dd36237 
[ghstack-poisoned]
@vkuzo vkuzo changed the base branch from gh/vkuzo/91/head to main June 24, 2025 19:27
vkuzo added a commit that referenced this pull request Jun 24, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
48ab874 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: ddb1f80
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo
Update
dc0803c 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Jun 25, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
3e27fca 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: b229983
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo vkuzo changed the title [wip] sharding strategy for dim1 kernel mxfp8 training: add TP sharding strategy for dim1 kernel Jun 25, 2025
vkuzo
vkuzo commented Jun 25, 2025
View reviewed changes
test/prototype/mx_formats/test_mx_dtensor.py
)
# TODO(future PR): enable compile here, currently seeing
# https://www.internalfb.com/phabricator/paste/view/P1851219639
# _test_lowp_mlp_tensor_parallelism_base(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need to uncomment this and run ./test/prototype/mx_formats/test_mx_dtensor.sh to reproduce

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @bdhirsh

@vkuzo vkuzo added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Jun 25, 2025
@vkuzo vkuzo requested review from drisspg and danielvegamyhre June 27, 2025 13:56
danielvegamyhre
danielvegamyhre approved these changes Jul 7, 2025
View reviewed changes
@vkuzo
Update
68b683d 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Jul 8, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
570770d 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: f66d58d
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo vkuzo mentioned this pull request Jul 8, 2025
Merged
@vkuzo vkuzo merged commit e675ffd into main Jul 9, 2025
47 of 53 checks passed
@jithunnair-amd
Copy link
Collaborator

@vkuzo This PR seems to have introduced timeouts in torch ao CI on trunk for both CUDA and ROCm nightly runs: https://hud.pytorch.org/hud/pytorch/ao/e675ffd9745e745056cd27a5f64cacad0aebd051/1?per_page=50&name_filter=regression&mergeEphemeralLF=true
cc @petrex @huydhn @drisspg

@danielvegamyhre
Copy link
Contributor

@vkuzo This PR seems to have introduced timeouts in torch ao CI on trunk for both CUDA and ROCm nightly runs: https://hud.pytorch.org/hud/pytorch/ao/e675ffd9745e745056cd27a5f64cacad0aebd051/1?per_page=50&name_filter=regression&mergeEphemeralLF=true cc @petrex @huydhn @drisspg

Just fyi we are increasing timeout thresholds in #2549 and #2548

@petrex
Copy link
Collaborator

petrex commented Jul 15, 2025

@vkuzo This PR seems to have introduced timeouts in torch ao CI on trunk for both CUDA and ROCm nightly runs: https://hud.pytorch.org/hud/pytorch/ao/e675ffd9745e745056cd27a5f64cacad0aebd051/1?per_page=50&name_filter=regression&mergeEphemeralLF=true cc @petrex @huydhn @drisspg

Just fyi we are increasing timeout thresholds in #2549 and #2548

thanks I've labeled CI flow for rocm for #2549 and #2548.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@danielvegamyhre danielvegamyhre danielvegamyhre approved these changes

@drisspg drisspg Awaiting requested review from drisspg

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@vkuzo @jithunnair-amd @danielvegamyhre @petrex @facebook-github-bot